from collections import defaultdict, deque
from typing import Dict, List
import os
import pickle
import numpy as np
from numpy.core.fromnumeric import mean
from numpy.lib.function_base import select
from pandas.core.frame import DataFrame
from src.datahelper import StockDataHelper
from src.backTestingData import BackTestingData
from src.taskSignals import TaskSignals
from src.strategy.baseStrategy import BaseStrategy
# from src.trainModel import TrainModel


class MlStrategy(BaseStrategy):
    def __init__(self) -> None:
        super().__init__()

    def init_strategy(self, dataHelper: StockDataHelper,
                      backTestingData: BackTestingData, taskSignals: TaskSignals):
        self.taskSignals = taskSignals
        self.backTestingData = backTestingData
        self.data_by_date = dataHelper.get_start_end_dayData(
            backTestingData.data['start_date'], backTestingData.data['end_date'])
        self.trade_day_list = list(self.data_by_date.keys())
        self.service_charge = backTestingData.data['service_charge']
        self.position.clear()
        self.daily_return.clear()
        self.money = self.backTestingData.data['money']
        self.initial_money = self.money
        # 初始买入的股票
        self.buy_stocks(
            backTestingData.data['selected_stocks'], self.trade_day_list[0], 'open')
        self.update_portfolio_value(self.trade_day_list[0])

        # self.next_n_day = self.backTestingData.data['next_n_day']
        self.next_n_day = 3
        self.window_size = 2
        # self.window_size = self.backTestingData.data['window_size']
        self.trans_period = self.backTestingData.data['transPeriod']
        self.trans_day_set = self.get_trans_period_day_set()  # 所有调仓日的list

        self.model = self.load_model()

    def get_hold_stocks(self):
        return list(filter(lambda k: self.position[k] != 0, self.position.keys()))

    def get_trans_period_day_set(self):
        """
        :return: 返回所有调仓日list, 隔k个交易日调仓一次
        """
        res = set()
        trans_period = self.trans_period
        for i in range(trans_period - 1, len(self.trade_day_list), trans_period):
            res.add(self.trade_day_list[i])
        return res

    def is_valid(self, data):
        for n in data:
            if n is None or n != n:
                return False
        return True

    # 获取交易日前window_size的所有factor值
    def get_window_data(self, trade_day):
        dataframe = self.data_by_date[trade_day]

        trade_day_idx = self.trade_day_list.index(trade_day)
        factor_dict = {}
        x_data = None
        stock_list = []
        # label_dict = {}
        for i in range(len(dataframe)):
            stock = dataframe.index[i]
            factor = []
            for j in range(self.window_size):  # 过去window_size天的所有数据
                idx = trade_day_idx - self.window_size + j  # 不包含当天数据
                pre_day = self.trade_day_list[idx]
                pre_df = self.data_by_date[pre_day]

                for keyword in ['volume', 'money', 'WVAD', 'OBV']:
                    pre_df[keyword] = pre_df[keyword] / 1000000

                pre_factor = pre_df.loc[stock].tolist()
                factor += pre_factor

            if not self.is_valid(factor):
                continue
            # factor_dict[stock] = factor
            stock_list.append(stock)

            x = np.array(factor, dtype=float).reshape(-1, len(factor))

            if x_data is None:
                x_data = x
            else:
                x_data = np.concatenate((x_data, x), axis=0)
        return x_data, stock_list

    def load_model(self):
        model_name = "svm_win%d_next%d.md" % (
            self.window_size, self.next_n_day)
        model_path = os.path.join(
            os.getcwd(), "src", "strategy", "svm", model_name)
        if not os.path.exists(model_path):
            self.log("The model not exists, Please train the %d window_size %d next_n_day model !" % (
                self.window_size, self.next_n_day))
            raise
        with open(model_path, 'rb') as f:
            model = pickle.load(f)
        return model

    # 预测未来n天的涨跌
    def pred_n_nex_day(self, trade_day):
        x_data, stock_list = self.get_window_data(trade_day)
        my_model = self.model
        pred_y = my_model.predict(x_data)
        res_dict = {}
        for i in range(len(stock_list)):
            res_dict[stock_list[i]] = pred_y[i]
        return res_dict

    # 调仓日更新
    def update_trans_day_info(self, date):
        sell_stocks, buy_stocks = [], []

        res_dict = self.pred_n_nex_day(date)
        for stock, y in res_dict.items():
            if y > 0:
                buy_stocks.append(stock)

        # 持有股票全部卖出
        sell_stocks = self.get_hold_stocks()
        # 更新仓位
        self.sell_stocks(sell_stocks, date)
        self.buy_stocks(buy_stocks, date)

    # 交易日的更新
    def update_trade_day_info(self, date):
        pass

    def run(self):
        for i, date in enumerate(self.trade_day_list):
            last = self.portfolio_value + self.money

            if (date in self.trans_day_set) and (self.trade_day_list.index(date) >= self.window_size):
                self.update_trans_day_info(date)
                self.taskSignals.progressBarSignal.emit(
                    int(i / len(self.trade_day_list) * 100))
            else:
                self.update_trade_day_info(date)

            # if self.trade_day_list.index(date) >= self.window_size:
            #     self.update_trans_day_info(date)
            #     self.taskSignals.progressBarSignal.emit(
            #         int(i / len(self.trade_day_list) * 100))
            # else:
            #     self.update_trade_day_info(date)

            self.update_portfolio_value(date)
            tot = self.portfolio_value + self.money

            # 收益率
            self.daily_return.append((tot - last) / self.initial_money * 100)
            self.log(f'{date}当前总资金{tot}')

        self.log(f'最大回撤{StockDataHelper.max_drawdown(self.daily_return)}')
        self.taskSignals.progressBarSignal.emit(100)  # 然后，让进度条更新到100%
        self.updateBackTestingData()
